from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection
import numpy as np
import os
from dotenv import load_dotenv

from utils.EmbeddingUtil import get_text_embedding

# 加载环境变量
load_dotenv()

class MilvusClient:
    def __init__(self, host=None, port=None):
        """
        初始化 Milvus 连接
        :param host: Milvus 服务的主机地址
        :param port: Milvus 服务的端口号
        """
        try:
            self.host = host or os.getenv("MILVUS_HOST", "localhost")
            self.port = port or os.getenv("MILVUS_PORT", "19530")
            self.alias = os.getenv("MILVUS_ALIAS", "default")  # 修正别名来源

            connections.connect(
                alias=self.alias,
                host=self.host,
                port=self.port
            )
            print(f"成功连接到 Milvus 服务 [{self.host}:{self.port}]")
        except Exception as e:
            print(f"连接 Milvus 服务失败: {e}")

    def create_collection(self, collection_name, dim, primary_field_name="id",
                         vector_field_name="embedding",
                         business_fields=None):
        """
        创建 Milvus 集合 (增强版)
        :param collection_name: 集合名称
        :param dim: 向量维度
        :param primary_field_name: 主键字段名称
        :param vector_field_name: 向量字段名称
        :param business_fields: 业务字段配置，格式为 [{"name":字段名, "dtype":数据类型, "max_length":最大长度}, ...]
        :return: 创建的集合对象
        """
        # 基础字段
        fields = [
            FieldSchema(name=primary_field_name, dtype=DataType.INT64, is_primary=True, auto_id=True),
            FieldSchema(name=vector_field_name, dtype=DataType.FLOAT_VECTOR, dim=dim)
        ]

        # 添加业务字段
        if business_fields:
            for field in business_fields:
                field_params = {
                    "name": field["name"],
                    "dtype": field["dtype"],
                    "description": field.get("description", "")
                }
                # 如果是VARCHAR类型，添加max_length参数
                if field["dtype"] == DataType.VARCHAR:
                    field_params["max_length"] = field.get("max_length", 255)

                fields.append(FieldSchema(**field_params))

        schema = CollectionSchema(fields=fields, description="RAG collection")
        collection = Collection(name=collection_name, schema=schema, using=self.alias)
        print(f"成功创建集合: {collection_name}")
        return collection

    def insert_data(self, collection, vectors, business_data=None,
                   vector_field_name="embedding"):
        """
        向 Milvus 集合中插入数据 (增强版)
        :param collection: 集合对象
        :param vectors: 要插入的向量列表
        :param business_data: 业务数据列表，格式为 [{"field1":value1, ...}, ...]
        :param vector_field_name: 向量字段名称
        :return: 插入结果
        """
        if business_data and len(business_data) != len(vectors):
            raise ValueError("业务数据数量必须与向量数量一致")

        data = []
        for i, vec in enumerate(vectors):
            item = {vector_field_name: vec}
            if business_data:
                item.update(business_data[i])
            data.append(item)

        insert_result = collection.insert(data)
        collection.flush()
        print(f"成功插入 {len(vectors)} 条数据")
        return insert_result

    def get_collection(self, collection_name):
        """
        通过集合名称获取集合对象
        :param collection_name: 集合名称
        :return: 集合对象，如果集合不存在则返回 None
        """
        try:
            collection = Collection(name=collection_name)
            return collection
        except Exception as e:
            print(f"获取集合 {collection_name} 失败: {e}")
            return None

    def insert_data(self, collection, vector, business_data=None,
                   vector_field_name="embedding"):
        """
        向 Milvus 集合中插入单条数据
        :param collection: 集合对象
        :param vector: 要插入的单个向量
        :param business_data: 单条业务数据，格式为 {"field1":value1, ...}
        :param vector_field_name: 向量字段名称
        :return: 插入结果
        """
        data = {vector_field_name: vector}
        if business_data:
            data.update(business_data)

        insert_result = collection.insert([data])
        collection.flush()
        print(f"成功插入 1 条数据")
        return insert_result

    def create_index(self, collection, vector_field_name="embedding", index_type="IVF_FLAT", metric_type="L2"):
        """
        为 Milvus 集合创建索引
        :param collection: 集合对象
        :param vector_field_name: 向量字段名称
        :param index_type: 索引类型
        :param metric_type: 距离度量类型
        """
        index_params = {
            "index_type": index_type,
            "metric_type": metric_type,
            "params": {"nlist": 128}
        }
        collection.create_index(field_name=vector_field_name, index_params=index_params)
        collection.load()
        print(f"成功为集合 {collection.name} 创建索引")

    def search(self, collection, query_vectors, top_k=10, vector_field_name="embedding", metric_type="L2"):
        """在 Milvus 集合中进行向量搜索（增加向量格式校验）"""
        # 转换 numpy 数组为列表
        if isinstance(query_vectors, np.ndarray):
            query_vectors = query_vectors.tolist()

        # 确保是二维列表格式
        if not isinstance(query_vectors, list) or not all(isinstance(v, list) for v in query_vectors):
            query_vectors = [query_vectors]  # 包裹一维向量

        # 添加维度校验
        # 修改字段访问方式
        collection_dim = next(
            (field.dim for field in collection.schema.fields
             if field.name == vector_field_name),
            None
        )
        if not collection_dim:
            raise ValueError(f"集合中未找到向量字段: {vector_field_name}")

        for vec in query_vectors:
            if len(vec) != collection_dim:
                raise ValueError(f"向量维度错误: 期望 {collection_dim} 维，实际 {len(vec)} 维")

        search_params = {
            "metric_type": metric_type,
            "params": {"nprobe": 10}
        }

        try:
            results = collection.search(
                data=query_vectors,
                anns_field=vector_field_name,
                param=search_params,
                limit=top_k,
                output_fields=["item_code", "item_name"]
            )
            print(results)
            # 转换结果为结构化数据
            formatted_results = []
            for result in results:
                for hit in result:
                    formatted_results.append({
                        "item_code": hit.entity.get("item_code"),
                        "item_name": hit.entity.get("item_name"),
                        "distance": hit.distance
                    })
            return formatted_results

        except Exception as e:
            print(f"搜索失败: {str(e)}")
            return []


# 使用示例
if __name__ == "__main__":
    # 初始化客户端
    client = MilvusClient()

    res = client.search(
        collection=client.get_collection("government_items"),
        query_vectors=[get_text_embedding("食品生产")],
        top_k=5
    )
    print(res)


